#!/usr/bin/env python3
"""
PostgreSQL数据库配置脚本
"""

import os
import sys

def setup_postgresql_config():
    """配置PostgreSQL连接参数"""
    print("=== PostgreSQL数据库配置 ===")
    print("请输入你的PostgreSQL数据库连接信息：")
    
    # 获取配置信息
    host = input("数据库主机 (默认: localhost): ").strip() or "localhost"
    port = input("数据库端口 (默认: 5432): ").strip() or "5432"
    database = input("数据库名称 (默认: arbitration): ").strip() or "arbitration"
    username = input("用户名 (默认: postgres): ").strip() or "postgres"
    password = input("密码: ").strip()
    
    # 设置环境变量
    os.environ["DATABASE_TYPE"] = "postgresql"
    os.environ["POSTGRES_HOST"] = host
    os.environ["POSTGRES_PORT"] = port
    os.environ["POSTGRES_DATABASE"] = database
    os.environ["POSTGRES_USERNAME"] = username
    os.environ["POSTGRES_PASSWORD"] = password
    os.environ["POSTGRES_SSLMODE"] = "prefer"
    
    print(f"\n配置完成！")
    print(f"主机: {host}")
    print(f"端口: {port}")
    print(f"数据库: {database}")
    print(f"用户名: {username}")
    print(f"密码: {'*' * len(password)}")
    
    return {
        "host": host,
        "port": port,
        "database": database,
        "username": username,
        "password": password
    }

def create_env_file(config):
    """创建环境变量文件"""
    env_content = f"""# PostgreSQL数据库配置
DATABASE_TYPE=postgresql
POSTGRES_HOST={config['host']}
POSTGRES_PORT={config['port']}
POSTGRES_DATABASE={config['database']}
POSTGRES_USERNAME={config['username']}
POSTGRES_PASSWORD={config['password']}
POSTGRES_SSLMODE=prefer
"""
    
    with open(".env", "w", encoding="utf-8") as f:
        f.write(env_content)
    
    print("\n已创建 .env 配置文件")

def test_postgresql_connection():
    """测试PostgreSQL连接"""
    print("\n=== 测试PostgreSQL连接 ===")
    
    try:
        from database_manager import DatabaseManager
        
        db = DatabaseManager()
        conn = db.get_connection()
        
        print("✓ 成功连接到PostgreSQL数据库")
        
        # 测试查询
        cases = db.query_cases()
        print(f"✓ 数据库初始化完成，包含 {len(cases)} 个案件")
        
        conn.close()
        print("✓ PostgreSQL连接测试通过！")
        return True
        
    except Exception as e:
        print(f"✗ PostgreSQL连接失败: {e}")
        print("\n可能的解决方案：")
        print("1. 检查PostgreSQL服务是否启动")
        print("2. 检查数据库连接参数是否正确")
        print("3. 检查数据库是否存在，如不存在请先创建")
        print("4. 检查用户权限是否足够")
        return False

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "test":
        test_postgresql_connection()
    else:
        config = setup_postgresql_config()
        create_env_file(config)
        
        if input("\n是否测试连接? (y/N): ").lower() == 'y':
            test_postgresql_connection()
